"""
Copyright 2020 Huawei Technologies Co., Ltd

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import time
import ascendfaiss
import faiss
import numpy as np


def read_binary(fname):
    a = np.fromfile(fname, dtype='float32')
    d = a.shape[0]
    a = a.reshape(int(d/512), 512)
    return a


load_index_from_file = True
start_time = time.time()
if not load_index_from_file:
    xb = read_binary("/home/qiuyue/G42_100million/base_100million.ivecs")
    xt = xb[:15000000, ]
xq = read_binary("/home/qiuyue/G42_100million/query_440000.ivecs")
xq = xq[:200000, ]
gt = np.fromfile("/home/qiuyue/G42_100million/groundtruth0-20W.ivecs", dtype='int64').reshape(200000, 1)
print("[%.3f] Read data from G42" % (time.time() - start_time))

if not load_index_from_file:
    d_in = xb.shape[1]    # vector dims
    d = d_in // 4
    print(d_in)
    print(d)
    nb = xb.shape[0]      # database size
    nt = xt.shape[0]      # tain set size
    print(nb)

nq = xq.shape[0]      # query size
nlist = 16384         # L1 IVFList size

# set running devices
dev = ascendfaiss.IntVector()
for i in range(0, 4):
    dev.push_back(i)

if not load_index_from_file:
    config = ascendfaiss.AscendIndexIVFSQConfig(dev)
    # create ivfsq index
    ascend_index_ivfsq = ascendfaiss.AscendIndexIVFSQ(
        d, nlist, ascendfaiss.ScalarQuantizer.QT_8bit,
        ascendfaiss.METRIC_L2, True, config)
    ascend_index_ivfsq.verbose = True

    ascend_index_pretransform = ascendfaiss.AscendIndexPreTransform(ascend_index_ivfsq)
    ascend_index_pretransform.prependPCA(d_in, d, 0, True)

    # train
    ascend_index_pretransform.train(xt)
    print("Training done!")

    cpu_index_pretransform = ascendfaiss.index_ascend_to_cpu(ascend_index_pretransform)
    faiss.write_index(cpu_index_pretransform, "pt.afterTrain.faiss")

    # add database
    ascend_index_pretransform.reserveMemory(nb)
    ascend_index_pretransform.add(xb)
    print("Adding done!")

    cpu_index_pretransform = ascendfaiss.index_ascend_to_cpu(ascend_index_pretransform)
    faiss.write_index(cpu_index_pretransform, "pt.afterAdd.faiss")
else:
    # restore the index
    cpu_index_pretransform = faiss.read_index("pt.afterAdd.faiss")
    ascend_index_pretransform = ascendfaiss.index_cpu_to_ascend(dev, cpu_index_pretransform)
    ascend_index_pretransform.verbose = True

# search topk results
k = 1000
warmup = xq[:31, :]
test_performance = True
random_query = False

if test_performance:
    repeat = 100
    batchArray = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1000]
    nprobeArray = [64]
    start = 60000
else:
    repeat = 1
    batchArray = [10000]
    nprobeArray = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
    start = 60000

for batch in batchArray:
    print(">>>>>>> batch size: %d" % batch)
    for nprobe in nprobeArray:
        ascendfaiss.AscendParameterSpace().set_index_parameter(
            ascend_index_pretransform, 'nprobe', nprobe)
        tmpeD, tempI = ascend_index_pretransform.search(warmup, k)
        cost_time = 0
        recall = {1: 0, 10: 0, 100: 0, 1000: 0}
        for times in range(repeat):
            if random_query:
                select_query = np.random.randint(nq, size=(batch,))
            else:
                select_query = np.array(list(range(start, start + batch)))
            t0 = time.time()
            Distance, IndexResult = ascend_index_pretransform.search(xq[select_query], k)
            t1 = time.time()
            cost_time += (t1 - t0) * 1000.0 / batch

            i = 1
            while i <= k:
                recall[i] += (np.array(IndexResult[:, :i] == gt[select_query, :1]) + 0).sum() / float(batch)
                i *= 10

        cost_time /= float(repeat)
        recall[1] /= float(repeat)
        recall[10] /= float(repeat)
        recall[100] /= float(repeat)
        recall[1000] /= float(repeat)

        print("@%3d qps: %.3f, r@1: %.4f, r@10: %.4f, r@100: %.4f, r@1000: %.4f" %
              (nprobe, 1000.0 / cost_time, recall[1], recall[10], recall[100], recall[1000]))
